

mata
mata clear

void ks_function(todo, p, y, g, negH) {
    eta = p[1]; sigma = p[2]; ro = p[3]
    observed_data = st_data(., "reports")
    XL_ro = ceil(ro)
    
    // Calculate x_hat values using vectorized approach
    x_hat = J(6, 1, .)
    for (i = 0; i <= 5; i++) {
        x_hat[i+1] = i >= XL_ro ? eta/((6) * (eta - i + ro)) : 
                     1/(6) * (1 - (2*normal((ro - i)/sigma) - 1))
    }
    
    // Calculate theta constraint
    sum1 = sum2 = 0
    for (y_val = 0; y_val <= XL_ro - 1; y_val++) {
        sum1 = sum1 + ((ro - y_val) <= 0 ? 0 : (2*normal((ro - y_val)/sigma) - 1))
    }
    for (x_val = XL_ro; x_val <= 5; x_val++) {
        sum2 = sum2 + (x_val - ro) / (eta - x_val + ro)
    }
    theta = sum1 - sum2
    
    // Calculate MSE for all cells
    mse_total = 0
    for (roll = 0; roll <= 5; roll++) {
        cell_pred =  x_hat[roll+1] / 6
        mse_roll = 0
        for (j = 1; j <= 6; j++) {
            obs_idx = roll * 6 + j
            mse_roll = mse_roll + (100*(cell_pred - observed_data[obs_idx]))^2
        }
        mse_total = mse_total + mse_roll
    }
    
    // Add penalty for theta constraint
    y = mse_total + 100000 * theta^2
}

// Optimization
S1 = optimize_init()
optimize_init_technique(S1, "nr")
optimize_init_which(S1, "min")
optimize_init_evaluator(S1, &ks_function())
optimize_init_params(S1,(3.65,2.96,3.05))


x = optimize(S1)
p = optimize_result_value(S1)

// Calculate final theta
eta_opt = x[1]; sigma_opt = x[2]; ro_opt = x[3]
XL_ro_opt = ceil(ro_opt)
sum1_opt = sum2_opt = 0
for (y = 0; y <= XL_ro_opt - 1; y++) {
    sum1_opt = sum1_opt + ((ro_opt - y) <= 0 ? 0 : (2*normal((ro_opt - y)/sigma_opt) - 1))
}
for (x_val = XL_ro_opt; x_val <= 5; x_val++) {
    sum2_opt = sum2_opt + (x_val - ro_opt) / (eta_opt - x_val + ro_opt)
}
theta_final = sum1_opt - sum2_opt

// Calculate predictions for all 36 cells
x_hat_final = J(6, 1, .)
for (i = 0; i <= 5; i++) {
    x_hat_final[i+1] = i >= XL_ro_opt ? eta_opt/((6) * (eta_opt - i + ro_opt)) : 
                       1/(6) * (1 - (2*normal((ro_opt - i)/sigma_opt) - 1))
}

predictions = J(36, 1, .)
for (roll = 0; roll <= 5; roll++) {
    cell_pred =  x_hat_final[roll+1] / 6
    for (j = 1; j <= 6; j++) {
        predictions[roll * 6 + j] = cell_pred
    }
}

// Store results as scalars
st_numscalar("eta", x[1])
st_numscalar("sigma", x[2])
st_numscalar("ro", x[3])
st_numscalar("mse", p)
st_numscalar("theta", theta_final)

end


// Add prediction variable to main dataset
gen predicted_ks_$treat  = .
mata: st_store(., "predicted_ks_$treat", predictions)


